#!/usr/bin/env python3
import os
import rospy
from std_msgs.msg import Float32
from balloon_rl.msg import nx_from_mcu
from balloon_controller.policy_controller import PolicyController

thrust_msg = Float32()

load_run = 'Aug24_14-34-28_'
checkpoint = 'model_500.pt'
model_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), '../model', load_run, checkpoint)
TARGET_HEIGHT = 1.2  # m

controller = PolicyController(cfg=None,
                              desired_height=TARGET_HEIGHT,  # m
                              model_path=model_path,
                              device='cpu')

def receive_callback(data):
    # Policy update (obs -> action)
    rospy.loginfo('receive data:{},{},{},{}'.format(data.height, data.roll, data.pitch, data.accZ)) 

    controller.update(data.height, data.accZ)
    action = controller.get_action()
    thrust_msg.data = action

def main():
    # In ROS, nodes are uniquely named. If two nodes with the same
    # name are launched, the previous one is kicked off. The
    # anonymous=True flag means that rospy will choose a unique
    # name for our 'listener' node so that multiple listeners can
    # run simultaneously.
    rospy.init_node('policy_update', anonymous=True)

    rospy.loginfo("policy_inference_node created!")

    rospy.Subscriber('MCU_DATA_receive', nx_from_mcu, receive_callback)

    pub = rospy.Publisher('MCU_DATA_send', Float32, queue_size=10)

    rate = rospy.Rate(10) # 10 Hz

    while not rospy.is_shutdown():
        # rospy.loginfo(thrust_msg.data)
        pub.publish(thrust_msg)
        rate.sleep()

    # spin() simply keeps python from exiting until this node is stopped
    rospy.spin()

if __name__ == '__main__':
    main()
